from mysql_type import get_db_config, get_data_sql, DataStore, save_book_total_pages
from json_type import get_data_json
from sort_API.ancient_reading_mode import ancient_overlap_sort, kmeans_up_down_split, dbscan_clusters_lines
from sort_API.modern_reading_mode import modern_overlap_sort, get_translation_field, dbscan_sort, get_pinyin_field, \
    get_line_id
import mysql.connector
from config import SAVE_PATH, SAVE_JSON, SAVE_IMAGE, SAVE_MYSQL, DISPLAY_DBSCAN, DISPLAY_OVERLAP, DISPLAY_TRANSLATION, \
    DISPLAY_PINYIN, DISPLAY_KMEANS
from save_API import save_json_file, save_image_file
from visual_image import modern_plot_visualization_data, modern_dbscan_overlap_data


class DataFromMode:
    """"功能判断类：用来处理不同来源的数据"""
    json_mode = 0
    database_mode = 1


class SortType:
    """"功能判断类：用来判断阅读类型"""
    modern_reading_mode = 0  # 按行阅读，从左往右
    ancient_reading_mode = 1  # 按列阅读，从左往右
    other = 2  # 防死锁参数


def main():
    """"主要参数调节器"""
    # 工具参数
    folder_path = 'json_type\input'  # 输入文件夹路径
    save_path = SAVE_PATH
    data_from_mode = DataFromMode.json_mode  # 设置数据来源

    book_id = None
    if data_from_mode == DataFromMode.database_mode:
        book_id = input("book_id:")

    # 根据数据来源进行不同的处理
    data = []
    if data_from_mode is DataFromMode.json_mode:
        data = get_data_json(folder_path)
    else:
        # 获取数据库配置
        db_config = get_db_config()
        # 连接数据库
        connection = mysql.connector.connect(**db_config)
        if connection.is_connected():
            data = get_data_sql(connection, book_id)

    reading_mode = data[0]['layout_direction']
    # reading_mode =SortType.ancient_reading_mode 临时设置阅读模式
    # 以页为单位进行排序
    page_idx = 0
    for page_data in data:
        page_idx += 1
        if len(page_data['page_word']) == 1 and page_data['page_word'][0]['word'] == "无法检测到明确的文字区域":
            continue
        
        sorted_data = []
        figure = None
        if reading_mode is SortType.modern_reading_mode:
            # 聚类分栏
            sorted_data_dbscan = dbscan_sort(page_data)
            if DISPLAY_DBSCAN:
                figure = modern_plot_visualization_data(sorted_data_dbscan, DISPLAY_DBSCAN)

            # 先去除非汉字、按行排序
            sorted_data_overlap = modern_overlap_sort(sorted_data_dbscan)
            if DISPLAY_OVERLAP:
                figure = modern_dbscan_overlap_data(sorted_data_overlap, DISPLAY_OVERLAP)

            # 获取翻译字段
            sorted_data_translation = get_translation_field(sorted_data_overlap)
            if DISPLAY_TRANSLATION:
                figure = modern_dbscan_overlap_data(sorted_data_translation, DISPLAY_TRANSLATION)

            # 获取拼音字段
            folder_name = page_data['title']
            sorted_data_pinyin = get_pinyin_field(page_idx, sorted_data_translation, folder_name)
            if DISPLAY_PINYIN:
                figure = modern_dbscan_overlap_data(sorted_data_pinyin, DISPLAY_PINYIN)

            sorted_data_line = get_line_id(sorted_data_pinyin, reading_mode)

            # 最终的排序数据
            parts = page_data['page_image_path'].split('/')[:3]
            ocr_source_images = '/'.join(parts)
            sorted_data = {
                'book_id': page_data['book_id'],
                'title': page_data['title'],
                'author': page_data['author'],
                'page_id': page_data['page_id'],
                'page_number': page_idx,
                'total_lines': len(sorted_data_line),
                'layout_direction': page_data['layout_direction'],
                'reading_sequence': page_data['reading_sequence'],
                'ocr_source_images': ocr_source_images,
                'page_image_path': page_data['page_image_path'],
                'page_word': sorted_data_line
            }


        elif reading_mode is SortType.ancient_reading_mode:
            # K-means聚类分栏--> 区分上下栏
            sorted_data_kmeans = kmeans_up_down_split(page_data)
            if DISPLAY_KMEANS:
                figure = modern_plot_visualization_data(sorted_data_kmeans, DISPLAY_KMEANS)

            # 采取重叠率来获得列
            sorted_data_overlap = ancient_overlap_sort(sorted_data_kmeans, page_data['reading_sequence'])
            if DISPLAY_OVERLAP:
                figure = modern_dbscan_overlap_data(sorted_data_overlap, DISPLAY_OVERLAP)

            sorted_data_line = get_line_id(sorted_data_overlap, reading_mode)

            # 最终的排序数据
            parts = page_data['page_image_path'].split('/')[:3]
            ocr_source_images = '/'.join(parts)
            sorted_data = {
                'book_id': page_data['book_id'],
                'title': page_data['title'],
                'author': page_data['author'],
                'page_id': page_data['page_id'],
                'page_number': page_idx,
                'total_lines': len(sorted_data_line),
                'layout_direction': page_data['layout_direction'],
                'reading_sequence': page_data['reading_sequence'],
                'ocr_source_images': ocr_source_images,
                'page_image_path': page_data['page_image_path'],
                'page_word': sorted_data_line
            }

        # 如果排序的数据不为空，则进行插入数据或者生成文件
        if sorted_data:
            if SAVE_JSON:
                save_json_file(sorted_data, save_path)
            if SAVE_IMAGE:
                save_image_file(sorted_data, save_path, figure)
            if SAVE_MYSQL:
                data_store = DataStore(sorted_data)
                data_store.save_sql_data()

    # 给书籍添加total_pages
    if SAVE_MYSQL:
        book_id = data[0]['book_id']
        save_book_total_pages(page_idx, book_id)


if __name__ == '__main__':
    main()
